#include <cmath>
#include <map>
#include <vector>

using namespace std;

/*********************************************************
 * Classes and typedefs
 *********************************************************/
class TreePtr;
class Point;

typedef multimap<int, Point *> TreeMap;
typedef vector<Point *> TreeVec;
typedef TreeMap::iterator MapIter;
/*********************************************************/


/*********************************************************
 * TreePtr
 *   Linear tree representation.
 *   Points exposed via two structures:
 *   	stl::multimap<int, Point *>
 *   	stl::vector<Point *>
 *********************************************************/
class TreePtr {
	public:
		int N;					// # of points
		int depth;				// depth of tree
		int shift;				// shift into high bits of key
		int ppt;				// points per thread
		double h;				// morton bucket size
		TreeMap M;				// multimap access to points
		TreeVec V;				// vector access to points
		TreeVec vL;				// vector access to leaves only
		Point *y;				// array of original points

		TreePtr (); 			
		TreePtr ( int );

		void (*G) (double *y, double *d, int N, double *x, double *u, int M);

		inline Point PointFromKey ( int );
		void PreorderPrint ();
		void PreorderPrintVec ();
		void BuildVector ();
		void BuildLeafVector ();
		void Insert ( Point );
		void parComputeSubtreeSizes ();
		void parComputeEulerTour ();
		void ComputeAverageHelper ( Point * );
		void ComputeAverage ();	
};

  // Default constructor. Make empty map.
TreePtr::TreePtr () { 
	M = TreeMap(); 
}

  // Construct with given depth, calculate shift.
TreePtr::TreePtr ( int D ) { 
	M = TreeMap(); 
	depth = D; 
	shift = (int)ceil(log2(depth));
}
/*********************************************************/


/*********************************************************
 * Point 
 *   Represents point in space and node in the tree.
 *   Refers back to its containing tree via T.
 *********************************************************/
class Point {
	public:
		double x;			// coordinate in space
		double d;			// density
		double density;		// average subtree density
		double mass;		// center of mass
		long int morton;	// full morton ID
		long int nodeID;	// truncated morton ID within tree
		short int level;	// level within tree
		short int s;		// subtree size
		short int i;		// in-rank in Euler tour
		short int o;		// out-rank in Euler tour
		bool visited;		// visited flag
		bool isLeaf;		// leaf status
		TreePtr *T;			// pointer to containing tree

		Point ();
		Point ( double, double );

		Point *ptrCopy ();
		int Left (); 
		int Right ();
		int Parent ();
};	

  // Default constructor. Null everything.
Point::Point () {
	x = d = density = mass = 0.0;
	morton = -1; 
	nodeID = 0;
	level = 0;
	s = 0;
	i = 0;
	o = 0;
	isLeaf = 0;
	visited = 0; 
}

  // Initialize x and density, null everything else.
Point::Point ( double inX, double inD ) {
	x = inX; 
	d = inD;
   	density = mass = 0.0; 
	morton = -1; 
	nodeID = 0;
	level = 0;
	s = 0;
	i = 0;
	o = 0;
	isLeaf = 0;
	visited = 0; 
}
/*********************************************************/



/*********************************************************
 * TreePtr functions 
 *********************************************************/
  // Return a point given a map key
inline Point TreePtr::PointFromKey ( int p ) { 
	return *((*(M.find(p))).second); 
} // PointFromKey


  // Print in preorder using multimap interface
void TreePtr::PreorderPrint () {
	MapIter it;
   	Point *p;

	for (it = M.begin(); it != M.end(); it++) {
		p = (*it).second;
		if (shift < 0) cout << "BAD SHIFT" << endl;
		cout << p->level << "," << (p->nodeID >> shift);
		cout << "," << p->morton << "," << p->Parent() << " ";
	}
	cout << endl;
} // PreorderPrint


  // Print in preorder using vector interface
void TreePtr::PreorderPrintVec () {
	unsigned int i;
   	Point *p;

	for (i = 0; i < V.size(); i++) {
		p = V[i]; 
		if (shift < 0) cout << "BAD SHIFT" << endl;
		cout << p->level << ",";
		cout << (p->nodeID >> shift);
		cout << "," << p->morton << "   ";	
		if (p->isLeaf) {
			cout << endl;	
			if (i + 1 < V.size()) {
				cout.width(V[i+1]->level * 9);
				cout << " ";
			}
		}
			
	}
	cout << endl;
} // PreorderPrintVec


  // Build vector in same order as multimap
void TreePtr::BuildLeafVector () {
	MapIter it;
	TreeVec vec;

	  // reserve full capacity to minimize realloc
	vec.reserve(N);
	for (it = M.begin(); it != M.end(); it++) 
		if ((*it).second->isLeaf)
			vec.push_back((*it).second);
	vL = vec;
} // BuildVector


  // Build vector in same order as multimap
void TreePtr::BuildVector () {
	MapIter it;
	TreeVec vec;

	  // reserve full capacity to minimize realloc
	vec.reserve(M.size());
	for (it = M.begin(); it != M.end(); it++) 
		vec.push_back((*it).second);
	V = vec;
} // BuildVector


  // Insert point into multimap
void TreePtr::Insert( Point pt ) {
	MapIter it;
	Point p;
	long int mask;

	  // make sure point has reference to containing tree
	pt.T = this;

	  // nodeID not in tree yet, insert
	if (M.count(pt.nodeID + pt.level) == 0) {
		pt.isLeaf = 1;
		M.insert(pair<long int, Point *>(pt.nodeID + pt.level, pt.ptrCopy()));
		return;
	} 

	  // find node in tree
	it = M.find(pt.nodeID + pt.level);
	  
	  // extract point information
	p = *((*it).second);

	  // if a leaf, push down the tree
	if (p.isLeaf) {
		  // mark as non-leaf
		(*it).second->isLeaf = 0;
		(*it).second->morton = 0;

  		  // increase level
		p.level += 1;

		  // determine morton mask for next level
		if (depth - p.level < 0) cout << "BAD SHIFT" << endl;
		mask = 1 << (depth - p.level);

		if (shift < 0) cout << "BAD SHIFT" << endl;

		  // send left or right depending on mask
		if (p.morton & mask) 
			p.nodeID = p.nodeID | (mask << shift); // right
		
		  // recursive insert
		Insert( p );
	}

	  // increase level of original point
	pt.level += 1;

	if (depth - p.level < 0) cout << "BAD SHIFT" << endl;

	  // determine mask
	mask = 1 << (depth - pt.level);

	if (shift < 0) cout << "BAD SHIFT" << endl;

	  // determine right or left
	if (pt.morton & mask) 
		pt.nodeID = pt.nodeID | (mask << shift); // right

	  // recursive insert
	Insert( pt );
} // Insert


  // Compute subtree by moving through preorder truncated
  // 	morton IDs until mismatch
void TreePtr::parComputeSubtreeSizes () {
	unsigned int i, j;

	#pragma omp parallel for private(j)
	for (i = 0; i < V.size() - 1; i++) {
		for (j = i + 1; j < V.size(); j++) 
			if (V[i]->nodeID != V[j]->nodeID) break;
	
		V[i]->s = j - (i + 1);
	}
} // parComputeSubtreeSizes


  // Compute Euler tour in parallel via in- and out-ranks
  // 	using tree levels and subtree sizes
void TreePtr::parComputeEulerTour () {
	unsigned int i;

	#pragma omp parallel for
	for (i = 0; i < V.size(); i++) {
		V[i]->i = i*2 - V[i]->level;
		V[i]->o = V[i]->i + 2*V[i]->s;
	}
} // parComputeEulerTour


  // Recursively compute all average densities/masses in the tree
void TreePtr::ComputeAverageHelper ( Point *p ) {
	Point *parent;
	short int visited = 0;

	#pragma omp critical
	if (p->visited == 1)
		visited = 1;
	else 
		p->visited = 1;

	if (visited) return;

	if (p->isLeaf) {
		p->density = p->d;
		p->mass = p->x;
	}

	if (p->Parent() == -1)
		return;

	parent = (*(M.find(p->Parent()))).second;

//	#pragma omp critical
	{
		parent->density += p->density;
		parent->mass += p->mass * p->density;
	}

	ComputeAverageHelper( parent );
} // ComputeAverageHelper

 
  // Wrapper to recursive average computation 
void TreePtr::ComputeAverage () {
	unsigned int i;

//	#pragma omp parallel for schedule(static, ppt) 
	for (i = 0; i < vL.size(); i++) 
		ComputeAverageHelper( vL[i] );
	
//	#pragma omp parallel for schedule(static, ppt) 
	for (i = 0; i < V.size(); i++) 
		V[i]->mass /= V[i]->density;

} // ComputeAverage


/*********************************************************
 * Point functions 
 *********************************************************/
  // Return a pointer to a copy of this point
Point *Point::ptrCopy () {
	Point *pt = new Point();
	pt->x = x;
	pt->d = d;
	pt->density = density;
	pt->mass = mass;
	pt->morton = morton; 
	pt->nodeID = nodeID;
	pt->level = level;
	pt->isLeaf = isLeaf;
	pt->s = s; 
	pt->i = i; 
	pt->o = o; 
	pt->T = T;
	return pt;
} // ptrCopy

  // Return key of left child of this point
int Point::Left () {
	return level + 1 + nodeID;
} // Left

  // Return key of right child of this point
int Point::Right () {
	long int key;
	int lev = level + 1;

	if (T->depth - lev + T->shift < 0) cout << "BAD SHIFT" << endl;
	key = lev + (nodeID | (1 << (T->depth - lev + T->shift)));
	
	return key;
} // Right

  // Return key of parent of this point
int Point::Parent () {
	long int key;
	unsigned long int mask;
	int lev = level - 1;

	if (T->depth - lev + T->shift < 0) cout << "BAD SHIFT" << endl;
	mask = ~0 << (T->depth - lev + T->shift);
	nodeID = nodeID & mask;
	key = lev + nodeID;

	return key;
} // Right
